open Expr

type import_path = {
	mutable imp_used : bool;
	imp_path : type_path;
	imp_pos : pos;
}

type import_wild = {
	mutable wimp_used : bool;
	wimp_path : string list;
	wimp_pos : pos;
}

type imports = {
	paths : (string,import_path) Hashtbl.t;
	mutable wildcards : import_wild list;
}

type type_decl =
	| Void
	| Dyn
	| Class of class_context
	| Static of class_context
	| Function of type_decl list * type_decl
	| Package of string list

and class_field = {
	f_name : string;
	f_type : type_decl;
	f_static : static_flag;
	f_public : public_flag;
	f_pos : pos;
}

and class_context = {
	path : type_path;
	param : class_context option;
	name : string;
	file : string;
	native : bool;
	interface : bool;
	dynamic : bool;
	imports : imports;
	fields : (string,class_field) Hashtbl.t;
	statics : (string,class_field) Hashtbl.t;
	mutable super : class_context;
	mutable implements : class_context list;
	mutable constructor : class_field option;
}

type local = {
	lt : type_decl;
	lf : int;
}

type context = {
	class_path : string list;
	files : (string,signature list) Hashtbl.t;
	classes : (type_path,class_context) Hashtbl.t;
	in_static : bool;
	in_lambda : class_context option;
	in_constructor : bool;
	locals : (string,local) Hashtbl.t;
	mutable frame : int;
	mutable inumber : type_decl;
	mutable ibool : type_decl;
	mutable istring : type_decl;
	mutable returns : type_decl;
	mutable current : class_context;
	mutable curwith : type_decl option;
	finalizers : (unit -> unit) list ref;
}


type error_msg =
	| Class_not_found of type_path
	| Class_name_mistake of type_path
	| Cannot_unify of type_decl * type_decl
	| Custom of string

exception Error of error_msg * pos
exception File_not_found of string

let verbose = ref false
let strict_mode = ref false
let use_components = ref false
let local_inference = ref false
let warn_imports = ref false

let argv_pos = { pfile = "<argv>"; pmin = -1; pmax = -1 }

let error msg p = raise (Error (msg,p))

let rec s_type_decl = function
	| Void -> "Void"
	| Dyn -> "Any"
	| Class c -> s_type_path c.path
	| Static c -> "#" ^ s_type_path c.path
	| Function (args,r) -> "function (" ^ String.concat ", " (List.map s_type_decl args) ^ ") : " ^ s_type_decl r
	| Package l -> String.concat "." l

let error_msg = function
	| Class_not_found p -> "class not found : " ^ s_type_path p
	| Class_name_mistake p -> "class name mistake : should be " ^ s_type_path p
	| Cannot_unify (ta,tb) -> s_type_decl ta ^ " should be " ^ s_type_decl tb
	| Custom msg -> msg

let verbose_msg m =
	if !verbose then begin
		print_endline m;
		flush stdout;
	end

let load_class_ref = ref ((fun _ -> assert false) : context -> type_path -> pos -> class_context)

let t_object ctx = !load_class_ref ctx ([],"Object") null_pos
let t_array ctx = !load_class_ref ctx ([],"Array") null_pos

let rec is_super sup c =
	if c == sup then
		true
	else if c.super == c then
		false
	else
		is_super sup c.super

let is_number ctx = function
	| Class c when c == (match ctx.inumber with Class c2 -> c2 | _ -> assert false) -> true
	| _ -> false

let is_boolean ctx = function
	| Class c when c == (match ctx.ibool with Class c2 -> c2 | _ -> assert false) -> true
	| _ -> false

let is_string ctx = function
	| Class c when c == (match ctx.istring with Class c2 -> c2 | _ -> assert false) -> true
	| _ -> false

let resolve_path ctx p pos =
	match p with
	| (_ :: _) , _ -> !load_class_ref ctx p pos
	| [] , n ->
		let rec loop = function
			| [] -> 
				(try
					let imp = Hashtbl.find ctx.current.imports.paths n in
					let cl = !load_class_ref ctx imp.imp_path pos in
					imp.imp_used <- true;
					cl
				with
					Not_found -> !load_class_ref ctx p pos)
			| imp :: l ->
				try
					let cl = !load_class_ref ctx (imp.wimp_path,n) pos in
					imp.wimp_used <- true;
					cl
				with
					Error (Class_not_found p,_) when p = (imp.wimp_path,n) -> loop l
		in
		loop ctx.current.imports.wildcards

let rec is_function cl =
	match cl.path with
	| ([],"Function") -> true
	| _ ->
		if cl.super == cl then
			false 
		else
			is_function cl.super

(* check that ta >= tb *)
let rec unify ta tb p =
	match ta , tb with
	| Dyn , x | x , Dyn -> ()
	| Void , Void -> ()
	| Function (args1,r1) , Function (args2,r2) ->
		let rec loop a1 a2 = 
			match a1 , a2 with
			| x :: l1, y :: l2 -> unify x y p; loop l1 l2
			| _ , _ -> ()
		in
		loop args1 args2;
		unify r2 r1 p
	| Class cl1, Class cl2 ->
		let rec loop cl1 =
			if cl1 == cl2 || List.exists loop cl1.implements then
				true
			else if cl1.super == cl1 then
				false
			else
				loop cl1.super
		in
		if not (loop cl1) then 
			(match cl1.param , cl2.param with 
			| Some c1 , Some c2 when c1 == c2 -> ()
			| _ -> error (Cannot_unify (ta,tb)) p)
	| Function _, Class c
	| Static _, Class c when c.super == c -> () (* unify with Object *)
	| Static _ , Class cl
	| Class cl, Static _
	| Class cl, Function _
	| Function _ , Class cl	when is_function cl -> ()
	| _ , _ ->
		error (Cannot_unify (ta,tb)) p

let unify_array t1 t2 v p =
	match t2 with
	| Class { path = ([t],"Array") } when t.[0] = '#' ->
		(match fst v with
		| ENew ((EStatic ([],"Array"),_),[]) -> ()
		| EArrayDecl _ -> ()
		| _ -> unify t1 t2 p);
	| _ -> unify t1 t2 p

let rec tcommon ctx ta tb p =
	match ta , tb with
	| Void , Void -> Void
	| Void , _
	| _ , Void -> error (Cannot_unify (ta,tb)) p
	| Dyn , _ | _ , Dyn -> Dyn
	| Function _ , _ | Static _ , _ -> tcommon ctx (Class (!load_class_ref ctx ([],"Function") null_pos)) tb p
	| _ , Function _ | _ , Static _ -> tcommon ctx ta (Class (!load_class_ref ctx ([],"Function") null_pos)) p
	| Package _ , _  | _ , Package _  -> assert false
	| Class a , Class b ->
		let rec is_sub cl1 cl2 =
			if cl1 == cl2 || List.exists (is_sub cl1) cl2.implements then
				true
			else if cl2.super == cl2 then
				false
			else
				is_sub cl1 cl2.super
		in
		let rec parent cl1 cl2 =
			if is_sub cl2 cl1 then
				Some cl2
			else
				let rec loop = function
					| [] -> if cl2.super == cl2 then None else parent cl1 cl2.super
					| i :: l ->
						match parent cl1 i with
						| Some i -> Some i
						| None -> loop l
				in
				loop cl2.implements
		in
		let p1 = parent a b in
		let p2 = parent b a in
		match p1 , p2 with
		| None, None -> Class (t_object ctx)
		| Some a, None -> Class a
		| None, Some b -> Class b
		| Some a, Some b -> if is_sub a b then Class b else Class a

let t_opt ctx p = function
	| None -> if !strict_mode && not ctx.current.native then error (Custom "Type required in strict mode") p; Dyn
	| Some ([],"Void") -> Void
	| Some t -> Class (resolve_path ctx t p)

let rec has_return any (e,p) =
	let has_return = has_return any in
	match e with
	| EVars _ 
	| EFunction _ 
	| EBreak
	| EContinue
	| EVal _ ->
		false
	| EReturn None
		-> any
	| EBlock el ->
		List.exists has_return el
	| EFor (el,_,_,e) ->
		List.exists has_return (e::el)
	| EForIn (e1,_,e2) ->
		has_return e1 || has_return e2
	| EIf (_,e,eo) ->
		has_return e || (match eo with None -> false | Some e -> has_return e)
	| EWhile (_,e,_) ->
		has_return e
	| ESwitch (_,cases) ->
		List.exists (fun (_,e) -> has_return e) cases
	| ETry (e,cl,fo) ->			
		(has_return e) || List.exists (fun (_,_,e) -> has_return e) !cl || (match fo with None -> false | Some e -> has_return e)
	| EWith (_,e) ->
		has_return e
	| EReturn (Some _ ) ->
		true

let ret_opt ctx p f =
	match f.fexpr with
	| Some e when not (has_return false e) -> 
		(match f.ftype with
		| None | Some ([],"Void") -> Void
		| Some cp -> error (Custom ("Missing return of type " ^ s_type_path cp)) p)
	| _ -> t_opt ctx p f.ftype

let rec add_class_field ctx clctx fname stat pub get ft p =
	if pub = IsPrivate && clctx.interface then error (Custom "Private fields are not possible in interfaces") p;
	if stat = IsStatic && clctx.interface then error (Custom "Static fields are not possible in interfaces") p;
	let h = (match stat with IsStatic -> clctx.statics | IsMember -> clctx.fields) in
	let f = (try Some (Hashtbl.find h fname) with Not_found -> None) in
	match get with
	| Getter | Setter ->
		add_class_field ctx clctx ((if get = Getter then "__get__" else "__set__") ^ fname) stat pub Normal ft p;
		let t = (if get = Getter then 
					(match ft with Function (_,x) -> x | _ -> assert false)
				else
					(match ft with 
					| Function ([arg],r) -> 
						if r <> Void && r <> Dyn then error (Custom "Setter should not return any value") p;
						arg 
					| Function _ ->
						error (Custom "Setter can only have one parameter") p
					| _ -> assert false)
				)
		in
		let f = (match f with 
		| None -> 
			{
				f_name = fname;
				f_type = t;
				f_static = stat;
				f_public = pub;
				f_pos = p;
			}
		| Some f ->
			{
				f_name = fname;
				f_type = begin (try unify f.f_type t f.f_pos; unify t f.f_type p; t with Error (Cannot_unify _,_) when !use_components -> f.f_type) end;
				f_static = stat;
				f_public = (if pub <> f.f_public then error (Custom "Getter and setter have different public/private visibility") p else pub);
				f_pos = p;
			}
		) in
		Hashtbl.replace h fname f
	| Normal ->
		if f <> None || Hashtbl.mem (match stat with IsStatic -> clctx.fields | IsMember -> clctx.statics) fname then error (Custom ("Field redefiniton : " ^ fname)) p;
		Hashtbl.add h fname {
			f_name = fname;
			f_type = ft;
			f_static = stat;
			f_public = pub;
			f_pos = p;
		}

let is_dynamic = function
	| Dyn | Function _ | Package _ -> true
	| Void -> false
	| Static c | Class c -> c.dynamic

let add_finalizer ctx f =
	ctx.finalizers := f :: !(ctx.finalizers)

let no_void t p =
	if t = Void then error (Custom "Void where Object expected") p

let define_local ctx name t p =
	if Hashtbl.mem ctx.locals name then error (Custom ("Local variable redefinition : " ^ name)) p;
	Hashtbl.add ctx.locals name { lt = t; lf = ctx.frame }

let new_frame ctx =
	let f = ctx.frame in
	ctx.frame <- ctx.frame + 1;
	f

let clean_frame ctx f =
	ctx.frame <- f;
	Hashtbl.iter (fun n l ->
		if l.lf > f then Hashtbl.remove ctx.locals n;
	) ctx.locals

let rec resolve t fname =
	match t with
	| Void
	| Dyn
	| Function _ -> None
	| Package p ->
		Some {
			f_name = fname;
			f_type = Package (p @ [fname]);
			f_static = IsMember;
			f_public = IsPublic;
			f_pos = null_pos;
		}
	| Static c -> 
		(try Some (Hashtbl.find c.statics fname) with Not_found -> if c.super == c then None else resolve (Static c.super) fname)
	| Class c -> 
		try 
			Some (Hashtbl.find c.fields fname)
		with 
			Not_found -> 
				if c.super == c then
					None 
				else 
					resolve (Class c.super) fname

and type_ident ctx name e p =
	(* with lookup *)
	try
		match ctx.curwith with
		| None -> raise Not_found
		| Some t ->
			match t with
			| Void
			| Static _
			| Package _ -> assert false
			| Dyn -> set_eval e (EStatic (["__With"],name)); Dyn
			| Function _ -> set_eval e (EStatic (["__With"],name)); Dyn
			| Class c ->
				match resolve (Class c) name with
				| None -> raise Not_found
				| Some { f_public = IsPrivate } when not (is_super c ctx.current) -> error (Custom "Cannot access private field") p
				| Some f -> set_eval e (EStatic (["__With"],name)); f.f_type
	with Not_found ->
	(* local variable lookup *)
	try
		let l = Hashtbl.find ctx.locals name in
		l.lt
	with
		Not_found ->
	(* member variable lookup *)
	try
		if name = snd ctx.current.path then begin
			set_eval e (EStatic ctx.current.path);
			Static ctx.current
		end else
		let f = (match resolve (Class ctx.current) name with None -> raise Not_found | Some f -> f) in
		if ctx.in_static then error (Custom ("Cannot access member variable " ^ name ^" in static function")) p;
		set_eval e (EField ((EConst (Ident "this"),p),name));
		f.f_type
	with
		Not_found ->
			(* static variable lookup *)
			let rec loop c =
				try 
					Some (c , Hashtbl.find c.statics name)
				with 
					Not_found -> 
						if c.super == c then None else loop c.super
			in
			match loop ctx.current with
			| Some (c,f) ->
				set_eval e (EField ((EStatic c.path,p),name));
				f.f_type
			| None -> 
				match resolve (Static (!load_class_ref ctx ([],"TopLevel") null_pos)) name with
				| Some f -> 
					if f.f_public = IsPublic then set_eval e (EField ((EConst (Ident "_global"),p),name));
					f.f_type
				| None -> 
					if String.length name > 6 && String.sub name 0 6 = "_level" && (try int_of_string (String.sub name 6 (String.length name - 6)) >= 0 with _ -> false) then
						Class (!load_class_ref ctx ([],"MovieClip") null_pos)
					else
						Package [name]

let type_constant ctx c e p =
	match c with
	| Int _ | Float _ -> ctx.inumber
	| String _ -> ctx.istring
	| Ident "_root" -> Class (!load_class_ref ctx ([],"MovieClip") p)
	| Ident "true" | Ident "false" -> ctx.ibool
	| Ident "null" | Ident "undefined" | Ident "_global" -> Dyn
	| Ident "this" ->
		if ctx.in_lambda <> None then
			Dyn
		else begin
			if ctx.in_static then error (Custom "Cannot access this in static function") p;
			Class ctx.current
		end
	| Ident "super" ->
		if ctx.in_lambda <> None then
			Dyn
		else begin
			if ctx.in_static then error (Custom "Cannot access super in static function") p;
			Class ctx.current.super
		end
	| Ident name ->
		type_ident ctx name e p

let rec resolve_package ctx v (p : string list) pos =
	match p with
	| [] -> assert false
	| cname :: fields ->
		let rec access p = function
			| [] -> EStatic p
			| x :: l -> EField ((access p l , pos), x)
		in
		let rec search_package p =
			let rec loop acc = function
				| [] -> raise Exit
				| x :: l ->
					let cpath = List.rev l , x in
					try
						let cl = !load_class_ref ctx cpath pos in
						let vv = access cl.path (List.rev acc) in
						set_eval v vv;
						Static cl , acc
					with
						Error (Class_not_found p,_) when p = cpath -> 
							loop (x :: acc) l
			in
			let t , fields = loop [] (List.rev p) in
			let rec loop t = function
				| [] -> t
				| f :: l ->
					loop (type_field ctx t f pos) l
			in
			loop t fields
		in
		let rec loop = function
			| [] ->
				let rec last = function
					| x :: [] -> x
					| x :: l -> last l
					| [] -> assert false
				in
				(match last p with
				| x when String.length x > 0 && x.[0] >= 'A' && x.[0] <= 'Z' -> error (Custom ("Unknown class " ^ String.concat "." p)) pos
				| _ -> error (Custom ("Unknown variable " ^ List.hd p)) pos)
			| (p,use) :: l ->
				try 
					let r = search_package p in
					use();
					r
				with
					Exit -> loop l
		in
		let p2 , n , use = (try 
			let imp = Hashtbl.find ctx.current.imports.paths cname in
			fst imp.imp_path, snd imp.imp_path, (fun() -> imp.imp_used <- true)
		with 
			Not_found -> 
				[] , cname, (fun() -> ())
		) in
		loop ((p2 @ n :: fields , use) :: (List.map (fun imp -> imp.wimp_path @ cname :: fields, (fun() -> imp.wimp_used <- true)) ctx.current.imports.wildcards))

and type_field ctx t f p =
	match resolve t f with
	| None -> 
		if not (is_dynamic t) then error (Custom (s_type_decl (match t with Static c -> Class c | _ -> t) ^ " has no " ^ (match t with Static _ -> "static " | _ -> "") ^ "field " ^ f)) p;
		Dyn
	| Some f ->
		if f.f_public = IsPrivate then (match t with
			| Class c | Static c ->
				if not (is_super c ctx.current) then begin
					if (match ctx.in_lambda with 
						| None -> true 
						| Some cur -> not (is_super c cur)
					) then error (Custom ("Cannot access private field " ^ f.f_name)) p
				end;
			| _ -> ());
		f.f_type